/*
 * Copyright (C) 2025 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use crate::common::constants::ErrorCode;
use crate::{log_e, log_i, Vec};
extern crate alloc;
use crate::utils::parcel::Parcel;
use alloc::{collections::BTreeMap, vec};
use core::mem;

#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, PartialOrd, Ord)]
#[repr(i32)]
pub enum AttributeKey {
    AttrRoot = 100000,
    AttrResultCode = 100001,
    AttrSignature = 100004,
    AttrIdentifyMode = 100005,
    AttrTemplateId = 100006,
    AttrTemplateIdList = 100007,
    AttrRemainAttempts = 100009,
    AttrLockoutDuration = 100010,
    AttrScheduleId = 100014,
    AttrData = 100020,
    AttrPinSubType = 100021,
    AttrScheduleMode = 100022,
    AttrPropertyMode = 100023,
    AttrType = 100024,
    AttrCredentialId = 100025,
    AttrController = 100026,
    AttrCallerUidD = 100027,
    AttrResult = 100028,
    AttrCapabilityLevel = 100029,
    AttrAlgorithmInfo = 100030,
    AttrTimeStamp = 100031,
    AttrRootSecret = 100032,
    AttrAttrs = 100033,
    AttrExpiredSysTime = 100034,
    AttrExecutorMatcher = 100036,
    AttrUserId = 100041,
    AttrToken = 100042,
    AttrExecutorRole = 100043,
    AttrEsl = 100044,
    AttrVerifierUdid = 100045,
    AttrCollectorUdid = 100046,
    AttrLocalUdid = 100063,
    AttrPeerUdid = 100064,
    AttrPublicKey = 100065,
    AttrChallenge = 100066,
    AttrExecutorIndex = 100067,
    AttrOldRootSecret = 100068,
    AttrAuthPurpose = 100069,
    AttrLockStateAuthType = 100075,
    AttrPropertyModeV2 = 100076,
    AttrReEnrollFlag = 100077,
}

impl TryFrom<i32> for AttributeKey {
    type Error = ErrorCode;
    fn try_from(value: i32) -> Result<Self, ErrorCode> {
        match value {
            100000 => Ok(AttributeKey::AttrRoot),
            100001 => Ok(AttributeKey::AttrResultCode),
            100004 => Ok(AttributeKey::AttrSignature),
            100005 => Ok(AttributeKey::AttrIdentifyMode),
            100006 => Ok(AttributeKey::AttrTemplateId),
            100007 => Ok(AttributeKey::AttrTemplateIdList),
            100009 => Ok(AttributeKey::AttrRemainAttempts),
            100010 => Ok(AttributeKey::AttrLockoutDuration),
            100014 => Ok(AttributeKey::AttrScheduleId),
            100020 => Ok(AttributeKey::AttrData),
            100021 => Ok(AttributeKey::AttrPinSubType),
            100022 => Ok(AttributeKey::AttrScheduleMode),
            100023 => Ok(AttributeKey::AttrPropertyMode),
            100024 => Ok(AttributeKey::AttrType),
            100025 => Ok(AttributeKey::AttrCredentialId),
            100026 => Ok(AttributeKey::AttrController),
            100027 => Ok(AttributeKey::AttrCallerUidD),
            100028 => Ok(AttributeKey::AttrResult),
            100029 => Ok(AttributeKey::AttrCapabilityLevel),
            100030 => Ok(AttributeKey::AttrAlgorithmInfo),
            100031 => Ok(AttributeKey::AttrTimeStamp),
            100032 => Ok(AttributeKey::AttrRootSecret),
            100033 => Ok(AttributeKey::AttrAttrs),
            100034 => Ok(AttributeKey::AttrExpiredSysTime),
            100036 => Ok(AttributeKey::AttrExecutorMatcher),
            100041 => Ok(AttributeKey::AttrUserId),
            100042 => Ok(AttributeKey::AttrToken),
            100043 => Ok(AttributeKey::AttrExecutorRole),
            100044 => Ok(AttributeKey::AttrEsl),
            100045 => Ok(AttributeKey::AttrVerifierUdid),
            100046 => Ok(AttributeKey::AttrCollectorUdid),
            100063 => Ok(AttributeKey::AttrLocalUdid),
            100064 => Ok(AttributeKey::AttrPeerUdid),
            100065 => Ok(AttributeKey::AttrPublicKey),
            100066 => Ok(AttributeKey::AttrChallenge),
            100067 => Ok(AttributeKey::AttrExecutorIndex),
            100068 => Ok(AttributeKey::AttrOldRootSecret),
            100069 => Ok(AttributeKey::AttrAuthPurpose),
            _ => Err(ErrorCode::GeneralError),
        }
    }
}

pub const MAX_SUB_MSG_NUM: usize = 10;
pub const MAX_EXECUTOR_MSG_LEN: usize = 2048;

#[derive(Debug, Clone)]
pub struct Attribute {
    map: BTreeMap<AttributeKey, Vec<u8>>,
}

impl Attribute {
    pub fn new() -> Self {
        Attribute {
            map: BTreeMap::new(),
        }
    }

    pub fn try_from_bytes(msg: &[u8]) -> Result<Self, ErrorCode> {
        if msg.is_empty() {
            log_e!("msg is empty");
            return Err(ErrorCode::BadParam);
        }

        let mut attribute = Attribute::new();
        let mut parcel = Parcel::from(msg);

        while parcel.has_next() {
            let attr_key = parcel.read_i32_le()?;
            let length = parcel.read_u32_le()? as usize;

            let mut data = vec![0u8; length];
            parcel.read_bytes(&mut data)?;

            match AttributeKey::try_from(attr_key) {
                Ok(key) => {
                    attribute.map.insert(key, data);
                }
                Err(_) => {
                    log_i!(" attribute {} not defined, skip", attr_key);
                }
            }
        }

        Ok(attribute)
    }

    pub fn to_bytes(&self) -> Result<Vec<u8>, ErrorCode> {
        let mut parcel = Parcel::new();

        for (&key, data) in &self.map {
            parcel.write_i32_le(key as i32);
            parcel.write_u32_le(data.len() as u32);
            parcel.write_bytes(data);
        }

        Ok(parcel.into_vec())
    }

    pub fn get_u32(&self, key: AttributeKey) -> Result<u32, ErrorCode> {
        let data = self.get_u8_slice(key)?;
        if data.len() != mem::size_of::<u32>() {
            log_e!("Invalid data size");
            return Err(ErrorCode::GeneralError);
        }

        Ok(u32::from_le_bytes(data.try_into().unwrap()))
    }

    pub fn set_u32(&mut self, key: AttributeKey, value: u32) {
        self.map.insert(key, value.to_le_bytes().to_vec());
    }

    pub fn get_i32(&self, key: AttributeKey) -> Result<i32, ErrorCode> {
        let data = self.get_u8_slice(key)?;
        if data.len() != mem::size_of::<i32>() {
            log_e!("Invalid data size");
            return Err(ErrorCode::GeneralError);
        }

        Ok(i32::from_le_bytes(data.try_into().unwrap()))
    }

    pub fn set_i32(&mut self, key: AttributeKey, value: i32) {
        self.map.insert(key, value.to_le_bytes().to_vec());
    }

    pub fn get_u64(&self, key: AttributeKey) -> Result<u64, ErrorCode> {
        let data = self.get_u8_slice(key)?;
        if data.len() != mem::size_of::<u64>() {
            log_e!("Invalid data size");
            return Err(ErrorCode::GeneralError);
        }

        Ok(u64::from_le_bytes(data.try_into().unwrap()))
    }

    pub fn set_u64(&mut self, key: AttributeKey, value: u64) {
        self.map.insert(key, value.to_le_bytes().to_vec());
    }

    pub fn get_u8_slice(&self, key: AttributeKey) -> Result<&[u8], ErrorCode> {
        self.map.get(&key).map(|val| val.as_slice()).ok_or_else(|| {
            log_e!("Attribute is not set, key:{:?}", key);
            ErrorCode::GeneralError
        })
    }

    pub fn fill_u8_slice(&self, key: AttributeKey, buffer: &mut [u8]) -> Result<(), ErrorCode> {
        let value = self.map.get(&key).map(|val| val.as_slice()).ok_or_else(|| {
            log_e!("Attribute is not set, key:{:?}", key);
            ErrorCode::GeneralError
        })?;
        
        if value.len() != buffer.len() {
            log_e!("Attribute size is unexpected, expected {} actual {}", buffer.len(), value.len());
            return Err(ErrorCode::GeneralError);
        }
        
        buffer.copy_from_slice(value);
        Ok(())
    }

    pub fn get_u8_vec(&self, key: AttributeKey) -> Result<&Vec<u8>, ErrorCode> {
        self.map.get(&key).ok_or_else(|| {
            log_e!("Attribute is not set, key:{:?}", key);
            ErrorCode::GeneralError
        })
    }

    pub fn set_u8_slice(&mut self, key: AttributeKey, data: &[u8]) {
        self.map.insert(key, data.to_vec());
    }

    pub fn get_u64_vec(&self, key: AttributeKey) -> Result<Vec<u64>, ErrorCode> {
        let data = self.get_u8_slice(key)?;
        let chunks = data.chunks_exact(mem::size_of::<u64>());
        if chunks.remainder().len() != 0 {
            log_e!("u8 slice length {} is incorrect", data.len());
            return Err(ErrorCode::GeneralError);
        }

        let mut u64_vec = Vec::with_capacity(chunks.len());
        for chunk in chunks {
            u64_vec.push(u64::from_le_bytes(chunk.try_into().unwrap()));
        }
        Ok(u64_vec)
    }

    pub fn set_u64_slice(&mut self, key: AttributeKey, data: &[u64]) -> () {
        let mut bytes = Vec::with_capacity(data.len() * mem::size_of::<u64>());
        for &value in data {
            bytes.extend_from_slice(&value.to_le_bytes());
        }
        self.map.insert(key, bytes);
    }

    pub fn get_u8_vecs(&self, key: AttributeKey) -> Result<Vec<Vec<u8>>, ErrorCode> {
        let mut parcel = Parcel::from(self.map.get(&key).map(|val| val.as_slice()).ok_or_else(
            || {
                log_e!("Attribute is not set, key:{:?}", key);
                ErrorCode::GeneralError
            },
        )?);
        let mut u8_vecs = Vec::new();
        while parcel.has_next() {
            let length = parcel.read_u32_le()? as usize;
            let mut data = vec![0u8; length];
            parcel.read_bytes(&mut data)?;
            u8_vecs.push(data);
        }
        Ok(u8_vecs)
    }

    pub fn set_u8_slices(&mut self, key: AttributeKey, u8_slices: &[&[u8]]) -> () {
        let mut parcel = Parcel::new();
        for &slice in u8_slices {
            parcel.write_u32_le(slice.len() as u32);
            parcel.write_bytes(slice);
        }
        self.map.insert(key, parcel.into_vec());
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::common::constants::ErrorCode;
    
    #[test]
    fn test_create_empty_attribute() {
        let attribute = Attribute::new();
        assert!(attribute.map.is_empty());
    }

    #[test]
    fn test_attribute_u32() {
        let mut attribute = Attribute::new();
        attribute.set_u32(AttributeKey::AttrLockoutDuration, 6036);
        attribute.set_u32(AttributeKey::AttrPropertyMode, 5697);

        assert_eq!(attribute.get_u32(AttributeKey::AttrLockoutDuration).unwrap(), 6036);
        assert_eq!(attribute.get_u32(AttributeKey::AttrPropertyMode).unwrap(), 5697);
        assert!(attribute.get_u32(AttributeKey::AttrResultCode).is_err());
    }

    #[test]
    fn test_attribute_i32() {
        let mut attribute = Attribute::new();
        attribute.set_i32(AttributeKey::AttrResultCode, 6036);
        attribute.set_i32(AttributeKey::AttrRemainAttempts, 5697);

        assert_eq!(attribute.get_i32(AttributeKey::AttrResultCode).unwrap(), 6036);
        assert_eq!(attribute.get_i32(AttributeKey::AttrRemainAttempts).unwrap(), 5697);
        assert!(attribute.get_i32(AttributeKey::AttrLockoutDuration).is_err());
    }

    #[test]
    fn test_attribute_u64() {
        let mut attribute = Attribute::new();
        attribute.set_u64(AttributeKey::AttrTemplateId, 6036);
        attribute.set_u64(AttributeKey::AttrScheduleId, 5697);

        assert_eq!(attribute.get_u64(AttributeKey::AttrTemplateId).unwrap(), 6036);
        assert_eq!(attribute.get_u64(AttributeKey::AttrScheduleId).unwrap(), 5697);
        assert!(attribute.get_u64(AttributeKey::AttrResultCode).is_err());
    }

    #[test]
    fn test_attribute_u8_array() {
        let mut attribute = Attribute::new();
        let data = vec![1, 2, 3, 4, 5];
        
        attribute.set_u8_slice(AttributeKey::AttrSignature, &data);
        assert_eq!(attribute.get_u8_slice(AttributeKey::AttrSignature).unwrap(), data.as_slice());
        
        attribute.set_u8_slice(AttributeKey::AttrSignature, &[]);
        assert!(attribute.get_u8_slice(AttributeKey::AttrSignature).unwrap().is_empty());
        
        assert!(attribute.get_u8_slice(AttributeKey::AttrResultCode).is_err());
    }

    #[test]
    fn test_attribute_u64_array() {
        let mut attribute = Attribute::new();
        let data = vec![123, 456, 789, 101112];
        
        attribute.set_u64_slice(AttributeKey::AttrTemplateIdList, &data);
        assert_eq!(attribute.get_u64_vec(AttributeKey::AttrTemplateIdList).unwrap(), data);
        
        attribute.set_u64_slice(AttributeKey::AttrTemplateIdList, &[]);
        assert!(attribute.get_u64_vec(AttributeKey::AttrTemplateIdList).unwrap().is_empty());
        
        assert!(attribute.get_u64_vec(AttributeKey::AttrSignature).is_err());
    }

    #[test]
    fn test_serialization_roundtrip() {
        let mut attribute = Attribute::new();
        attribute.set_i32(AttributeKey::AttrResultCode, 123);
        attribute.set_u32(AttributeKey::AttrIdentifyMode, 456);
        attribute.set_u64(AttributeKey::AttrScheduleId, 789);
        attribute.set_u8_slice(AttributeKey::AttrSignature, &[1, 2, 3]);

        let serialized = attribute.to_bytes().unwrap();
        let deserialized = Attribute::try_from_bytes(&serialized).unwrap();

        assert_eq!(deserialized.get_i32(AttributeKey::AttrResultCode).unwrap(), 123);
        assert_eq!(deserialized.get_u32(AttributeKey::AttrIdentifyMode).unwrap(), 456);
        assert_eq!(deserialized.get_u64(AttributeKey::AttrScheduleId).unwrap(), 789);
        assert_eq!(deserialized.get_u8_slice(AttributeKey::AttrSignature).unwrap(), &[1, 2, 3]);
    }

    #[test]
    fn test_invalid_deserialization() {
        assert!(Attribute::try_from_bytes(&[]).is_err());
        
        let partial_data = vec![0x01, 0x00, 0x00, 0x00];
        assert!(Attribute::try_from_bytes(&partial_data).is_err());
        
        let mut invalid_data = vec![];
        invalid_data.extend(&(AttributeKey::AttrResultCode as i32).to_le_bytes());
        invalid_data.extend(&1000_u32.to_le_bytes());
        invalid_data.extend(&[1, 2, 3]);
        assert!(Attribute::try_from_bytes(&invalid_data).is_err());
    }

    #[test]
    fn test_complex_attribute_serialization() {
        let mut attribute = Attribute::new();
        let keys = [
            AttributeKey::AttrResultCode,
            AttributeKey::AttrSignature,
            AttributeKey::AttrData,
            AttributeKey::AttrRemainAttempts,
        ];

        for (i, &key) in keys.iter().enumerate() {
            let data_len = (i + 1) * 10;
            let data: Vec<u8> = (0..data_len).map(|j| (j % 256) as u8).collect();
            attribute.set_u8_slice(key, &data);
        }

        let serialized = attribute.to_bytes().unwrap();
        let deserialized = Attribute::try_from_bytes(&serialized).unwrap();

        for &key in &keys {
            let original = attribute.get_u8_slice(key).unwrap();
            let restored = deserialized.get_u8_slice(key).unwrap();
            assert_eq!(original, restored);
        }
    }

    #[test]
    fn test_unsupported_key_handling() {
        let mut data = vec![];
        
        data.extend(&0xFFFFFFFF_u32.to_le_bytes());
        data.extend(&4_u32.to_le_bytes());
        data.extend(&[1, 2, 3, 4]);
        
        let attribute = Attribute::try_from_bytes(&data).unwrap();
        assert!(attribute.map.is_empty());
    }

    #[test]
    fn test_empty_attribute_serialization() {
        let attribute = Attribute::new();
        let serialized = attribute.to_bytes().unwrap();
        assert!(serialized.is_empty());
    }
}